#include "ai_service_client.h"

#include <thread>
#include <securec.h>

#include "ai_service_death_recipient.h"
#include "hisysevent.h"
#include "iservice_registry.h"
#include "system_ability_definition.h"

#include "ai_service_log.h"
#include "ai_service_errors.h"

using namespace cv;

namespace OHOS {
namespace AIS {
namespace {
    constexpr OHOS::HiviewDFX::HiLogLabel LABEL = {LOG_CORE, AIS_LOG_DOMAIN, "AIServiceClient"};
    constexpr int32_t GET_SERVICE_MAX_COUNT = 30;
    constexpr uint32_t WAIT_MS = 200;
    constexpr int INPUT_WIDTH = 640;
    constexpr int INPUT_HEIGHT = 640;
} // namespace

int32_t AIServiceClient::InitAIServiceClient()
{
    std::lock_guard<std::mutex> clientLock(clientMutex_);
    if (aiServiceProxy_ != nullptr)
    {
        AIS_HILOGD("aiServiceProxy_ already init");
        return RETCODE_SUCCESS;
    }
    auto systemManager = SystemAbilityManagerClient::GetInstance().GetSystemAbilityManager();
    CHKPTR(systemManager, RETCODE_GET_SA_MANAGER_FAILED);
    int32_t retry = 0;
    while (retry < GET_SERVICE_MAX_COUNT) {
        aiServiceProxy_ = iface_cast<IAIService>(systemManager->GetSystemAbility(
            AI_SERVICE_ID));
        if (aiServiceProxy_ != nullptr)
        {
            AIS_HILOGD("aiServiceProxy_ get service success, retry:%{public}d", retry);
            serviceDeathObserver_ = new (std::nothrow) AIServiceDeathRecipient(*const_cast<AIServiceClient *>(this));
            if (serviceDeathObserver_ != nullptr)
            {
                aiServiceProxy_->AsObject()->AddDeathRecipient(serviceDeathObserver_);
            }
            return RETCODE_SUCCESS;
        }
        AIS_HILOGW("Get service failed, retry:%{public}d", retry);
        std::this_thread::sleep_for(std::chrono::milliseconds(WAIT_MS));
        retry++;
    }
    AIS_HILOGE("Get service failed");
    return RETCODE_INIT_CLIENT_FAILED;
}

void AIServiceClient::ProcessDeathObserver(wptr<IRemoteObject> object)
{
    (void)object;
    aiServiceProxy_ = nullptr;
    auto ret = InitAIServiceClient();
    if (ret != RETCODE_SUCCESS)
    {
        AIS_HILOGE("InitAIServiceClient failed, ret:%{public}d", ret);
        return;
    }
}

int32_t AIServiceClient::Init(std::string& modelPath, int32_t* modelId)
{
    auto ret = InitAIServiceClient();
    if (ret != RETCODE_SUCCESS)
    {
        AIS_HILOGE("InitAIServiceClient failed, ret : %{public}d", ret);
        return ret;
    }
    CHKPTR(aiServiceProxy_, RETCODE_NULL_PARAM);
    return aiServiceProxy_->Init(modelPath, modelId);
}

int32_t AIServiceClient::Process(const int32_t modelId, PicDesc& picDesc, std::string& picData, InferResult& result)
{
    auto ret = CheckInput(picDesc, picData);
    if (ret != RETCODE_SUCCESS)
    {
        return ret;
    }

    std::string resizeBuffer = "";
    ret = PreProcess(picDesc, picData, resizeBuffer);
    if (ret != RETCODE_SUCCESS)
    {
        return ret;
    }

    ret = InitAIServiceClient();
    if (ret != RETCODE_SUCCESS)
    {
        AIS_HILOGE("InitAIServiceClient failed, ret : %{public}d", ret);
        return ret;
    }
    CHKPTR(aiServiceProxy_, RETCODE_NULL_PARAM);

    return aiServiceProxy_->Process(modelId, picDesc, resizeBuffer, result);
}

int32_t AIServiceClient::DeInit()
{
    auto ret = InitAIServiceClient();
    if (ret != RETCODE_SUCCESS)
    {
        AIS_HILOGE("InitAIServiceClient failed, ret : %{public}d", ret);
        return ret;
    }
    CHKPTR(aiServiceProxy_, RETCODE_NULL_PARAM);
    return aiServiceProxy_->DeInit();
}

int32_t AIServiceClient::CheckInput(PicDesc& picDesc, std::string inputData) const
{
    if (picDesc.width == 0 || picDesc.height == 0)
    {
        AIS_HILOGE("width and height in picDesc is invalid.");
        return RETCODE_INVALID_PARAM;
    }
    
    if (picDesc.dataSize == 0 || picDesc.dataSize != inputData.size())
    {
        AIS_HILOGE("dataSize in picDesc is zero or not equal to input data size.");
        return RETCODE_INVALID_PARAM;
    }

    return RETCODE_SUCCESS;   
}

int32_t AIServiceClient::PreProcess(PicDesc& picDesc, std::string& picData, std::string& outBuff)
{
    AIS_HILOGD("PreProcess start.");

    if (picData.empty()) {
        AIS_HILOGE("PreProcess: input pic data is empty.");
        return RETCODE_INVALID_PARAM;
    }
    AIS_HILOGD("picData size: %{public}d.", (int)picData.size());
    std::vector<unsigned char> inBuffer(picData.begin(), picData.end());
    Mat img = imdecode(inBuffer, 1);
    if (img.data == nullptr || img.channels() != 3 || img.rows == 0 || img.cols == 0) {
        AIS_HILOGE("PreProcess: imdecode failed.");
        return RETCODE_OPENCV_IMDECODE_FAILED;
    }

    LetterBox letterBox {};
    letterBox.targetHeight = INPUT_HEIGHT;
    letterBox.targetWidth = INPUT_WIDTH;
    letterBox.inHeight = img.rows;
    letterBox.inWidth = img.cols;

    Mat blob;
    if ((letterBox.inHeight != letterBox.targetHeight) || (letterBox.inWidth != letterBox.targetWidth))
    {
        ComputeLetterBoxAndResize(letterBox, img, blob);
        letterBox.reverseAvailable = true;
    } else {
        blob = img;
    }
    picDesc.letterBox = letterBox;
    std::vector<unsigned char> resizeBuff;
    imencode(".jpg", blob, resizeBuff);
    outBuff.resize(resizeBuff.size());
    memcpy(&outBuff[0], resizeBuff.data(), resizeBuff.size());

    return RETCODE_SUCCESS;
}

void AIServiceClient::ComputeLetterBoxAndResize(LetterBox& lb, cv::Mat &src, cv::Mat &dest)
{
    lb.imgWHratio = (float)lb.inWidth / (float)lb.inHeight;
    lb.targetWHratio = (float)lb.targetWidth / (float)lb.targetHeight;
    if (lb.imgWHratio >= lb.targetWHratio) {
        HeightPad(lb, src, dest);
    } else {
        WidthPad(lb, src, dest);
    }
}

void AIServiceClient::HeightPad(LetterBox& lb, cv::Mat &src, cv::Mat &dest)
{
    // pad height dim
    lb.resizeScaleW = (float)lb.targetWidth / (float)lb.inWidth;
    lb.resizeScaleH = lb.resizeScaleW;
    lb.resizeWidth = lb.targetWidth;
    lb.wPadLeft = 0;
    lb.wPadRight = 0;

    lb.resizeHeight = (int)((float)lb.inHeight * lb.resizeScaleH);
    lb.hPadTop = (lb.targetHeight - lb.resizeHeight) / 2;

    if (((lb.targetHeight - lb.resizeHeight) % 2) == 0) {
        lb.hPadBottom = lb.hPadTop;
    } else {
        lb.hPadBottom = lb.hPadTop + 1;
    }

    dest = Mat(lb.targetHeight, lb.targetWidth, CV_8UC3, Scalar(0, 0, 0));
    Mat resizeImg;
    resize(src, resizeImg, Size(lb.resizeWidth - 1, lb.resizeHeight - 1), 0, 0, INTER_LINEAR);
    int h_off = int((lb.targetHeight - lb.resizeHeight) / 2);
    Mat roiImage = dest(Rect(0, h_off, lb.resizeWidth - 1, lb.resizeHeight - 1));
    resizeImg.copyTo(roiImage, resizeImg);
}

void AIServiceClient::WidthPad(LetterBox& lb, cv::Mat &src, cv::Mat &dest)
{
    // pad width dim
    lb.resizeScaleH = (float)lb.targetHeight / (float)lb.inHeight;
    lb.resizeScaleW = lb.resizeScaleH;

    lb.resizeWidth = (int)((float)lb.inWidth * lb.resizeScaleW);
    lb.wPadLeft = (lb.targetWidth - lb.resizeWidth) / 2;
    if (((lb.targetWidth - lb.resizeWidth) % 2) == 0) {
        lb.wPadRight = lb.wPadLeft;
    } else {
        lb.wPadRight = lb.wPadLeft + 1;
    }

    lb.resizeHeight = lb.targetHeight;
    lb.hPadTop = 0;
    lb.hPadBottom = 0;

    dest = Mat(lb.targetHeight, lb.targetWidth, CV_8UC3, Scalar(0, 0, 0));
    Mat resizeImg;
    resize(src, resizeImg, Size(lb.resizeWidth, lb.resizeHeight), 0, 0, INTER_LINEAR);
    int w_off = int((lb.targetWidth - lb.resizeWidth) / 2);
    Mat roiImage = dest(Rect(w_off, 0, lb.resizeWidth, lb.resizeHeight));
    resizeImg.copyTo(roiImage, resizeImg);
}

}  // namespace AIS    
}  // namespace OHOS